import { AObject } from "./AObject";
import { AZone } from "./AZone";
import { Agent } from "./Agent";
import * as Debug from "./Debug";
import { NAlgorithms } from "./NarrativeAlgorithms";
import { N, NUtils, PluralPronouns, Pronoun, PronounList, PronounNames, ProtectedPronouns, SP, Sentence, SentenceFlags, Signifier, SyntaxPattern, VerbElement, type Referent, type SyntaxElement, type SyntaxNode } from "./NarrativeBase";
import { U } from "./Util";

class CombiningOptimizer {
    constructor(
        public pattern: SyntaxPattern[],
        public tryCombine: (s0: Sentence, s: Sentence, m0: Map<string, SyntaxElement[]>, 
            m: Map<string, SyntaxElement[]>) => boolean) 
    {};
    
    process(narrative: Sentence[]): Sentence[] {
        let newNarrative: Sentence[] = [];
        let s0: Sentence | undefined;
        while (s0 = narrative.shift()) {
            let map0 = SyntaxPattern.match(this.pattern, s0.elements);
            if (map0) {
                // unmergeable sentences go here
                let clean: Sentence[] = [];
                let i = 0;
                // merge all combinable sentences into the first
                for (; i < narrative.length; i++) {
                    let s = narrative[i];
                    // stop when group changes
                    if (s.group != s0.group) break;
                    let map = SyntaxPattern.match(this.pattern, s.elements);
                    if (!map || !this.tryCombine(s0, s, map0, map))
                        clean.push(s);
                }
                // the rest is also unmergeable
                clean.push(...narrative.slice(i));
                narrative = clean;
            }
            newNarrative.push(s0);
        }
        return newNarrative;
    }
}

class _Cache {
    [Pronoun.You]?: Referent;
    [Pronoun.I]?:  Referent;
    [Pronoun.She]?: Referent;
    [Pronoun.He]?:  Referent;
    [Pronoun.It]?:  Referent;
    [Pronoun.TheyInanimate]?: Referent[];
};

export class NarrativeContext {
    cache = new _Cache();
    freeIndirectSubject = new Set<Referent>;
    proximatlySetPronouns = new Set<Pronoun>;
    _debug = false;

    setPronoun<P extends keyof _Cache>(p: P, r: _Cache[P]) {
        this.cache[p] = r;
    }

    queryPronoun<P extends keyof _Cache>(p: P): _Cache[P] {
        return this.cache[p];
    }

    clone(): NarrativeContext {
        const result = new NarrativeContext;
        result.freeIndirectSubject = new Set(this.freeIndirectSubject);
        result.proximatlySetPronouns = new Set(this.proximatlySetPronouns);
        Object.assign(result.cache, this.cache);
        if (this.cache[Pronoun.TheyInanimate])
            result.cache[Pronoun.TheyInanimate] = 
                [...this.cache[Pronoun.TheyInanimate]];
        return result;
    }

    static #getPronoun(elem: SyntaxElement): [Pronoun | null, Referent[] | Referent] | [null, null] {
        if (elem instanceof Signifier) {
            const pronoun = elem.object.pronoun;
            return [pronoun, PluralPronouns.includes(pronoun as any) ? 
                [elem.object] : elem.object];
        } else if (NUtils.isNode(elem) && elem.type == 'nounPhraseList') {
            if (!elem.children.every((x) => x instanceof Signifier)) {
                Debug.warn('NP-LIST currently allows only signifiers');
                return [null, null];
            }
            const list = elem.children.map((x) => (x as Signifier).object)
            if (new Set(list).size != elem.children.length)
                Debug.warn('NP-LIST has duplicate items');
            return [Pronoun.TheyInanimate, list];
        };
        return [null, null];
    }

    processElementPronoun(elem: SyntaxElement) {
        // is the element referable by a pronoun?
        const [pronoun, object] = NarrativeContext.#getPronoun(elem);
        if (pronoun === null) return;
        Debug.assert(object !== null); // TS type inference is lame
        // employ the pronoun if suitable
        if (Array.isArray(object)) {
            if (this.cache[pronoun] && (this.cache[pronoun] as Referent[]).every(
                (x) => (object as Referent[]).includes(x))) return pronoun;
        } else if (this.cache[pronoun] == object) return pronoun;

        // it is something newly referred to.
        // then, first check to not modify pronouns such as 'I' and 'you'
        if (ProtectedPronouns.includes(pronoun)) return;
        // if two referents that are represented by the same pronoun are too close, unset 
        // that pronoun instead of setting it
        if (this.proximatlySetPronouns.has(pronoun) && this.cache[pronoun])
            this.cache[pronoun] = undefined;
        else {
            // should unset 3rd person plural when setting the singular and vice versa
            if (pronoun == Pronoun.TheyInanimate) {
                this.cache[pronoun] = object as Referent[];
                this.cache[Pronoun.It] = undefined;
            } else {
                this.cache[pronoun] = object as Referent;
                this.cache[Pronoun.TheyInanimate] = undefined;
            }
            if (this._debug) console.log(Pronoun[pronoun], '->', [object].flat().map(
                (x) => typeof x.name == 'string' ? x.name : x.name.value));
            this.proximatlySetPronouns.add(pronoun);
        }
    }

    processSentencePronoun(orig: SyntaxElement[]): SyntaxElement[] {
        const list = [...orig];
        for (let i = 0; i < list.length; i++) {
            const elem = list[i];
            if (NUtils.isNode(elem)) {
                if (elem.type == 'coordinatingPredicatives') {
                    elem.children = elem.children.map((x) => {
                        let result = this.processSentencePronoun([x])[0];
                        this.proximatlySetPronouns.clear();
                        return result;
                    });
                } else if (elem.type != 'nounPhraseList') {
                    elem.children = this.processSentencePronoun(elem.children);
                    continue;
                }
            }
            const pronoun = this.processElementPronoun(elem);
            if (pronoun !== undefined) list[i] = PronounNames[pronoun];
        }
        this.proximatlySetPronouns.clear();
        return list;
    }
}

export type ConfabulationCallback = (elem: SyntaxElement, str: string) => string | void;

export class Narrative {
    context = new NarrativeContext;

    /** 
     * Combine simple sentences with same verb and object.
     * A摸了摸X。B摸了摸X。=> A和B摸了摸X。
     */
    static #Opt1 = new CombiningOptimizer(
        [SP.S('subj'), SP.P().with(SP.any('any').asterisk(), SP.O('obj'))],
        (s0, s, m0, m) => 
    {
        if (!NUtils.allEquals(m0.get('obj'), m.get('obj'))) return false;
        if (!NUtils.allEquals(m0.get('any'), m.get('any'))) return false;
        const subj0 = m0.get('subj')![0] as SyntaxNode;
        const subj = m.get('subj')![0] as SyntaxNode;
        subj0.children = [NUtils.combineElements(...subj0.children, ...subj.children)];
        return true;
    });

    /** 
     * Combine simple sentences with same subject and verb.
     * A拿起X。A拿起Y。=> A拿起X和Y。
     */
    static #Opt2 = new CombiningOptimizer(
        [SP.S('subj'), SP.P().with(SP.any('any').asterisk(), SP.O('obj'))],
        (s0, s, m0, m) => 
    {
        if (!NUtils.allEquals(m0.get('subj'), m.get('subj'))) return false;
        if (!NUtils.allEquals(m0.get('any'), m.get('any'))) return false;
        const obj0 = m0.get('obj')![0] as SyntaxNode;
        const obj = m.get('obj')![0] as SyntaxNode;
        obj0.children = [NUtils.combineElements(...obj0.children, ...obj.children)];
        return true;
    });

    /**
     * Tweaks structure according to pronouns that will be used.
     * - A拿起X和Y。A吃了X。=> A拿起X。A拿起Y。A吃了Y。(in Opt4, the last Y becomes 它)
     * - A摸了摸Z。A拿起X、Y和Z。=> A摸了摸Z。A拿起Z。A拿起X和Y。(in Opt4, the Z in the middle becomes 它)
     */
    #Opt3(narrative: Sentence[]) {
        function removeIfInList(
            node: SyntaxNode, ref: Referent | Referent[] | undefined)
        {
            if (!ref) return null;
            if (Array.isArray(ref)) {
                // assuming no repeated references in node...
                if (ref.length > node.children.length) return null;
                const newChildren: SyntaxElement[] = [...node.children];
                const separated: Signifier[] = [];
                for (let x of ref) {
                    let index = newChildren.findIndex(
                        (y) => y instanceof Signifier && y.object == x);
                    if (index < 0) return null;
                    separated.push(...newChildren.splice(index, 1) as Signifier[]);
                }
                node.children = newChildren;
                return separated;
            } else for (let i = 0; i < node.children.length; i++) {
                let elem = node.children[i];
                if (elem instanceof Signifier && elem.object == ref)
                    return node.children.splice(i, 1)[0] as Signifier;
            }
            return null;
        }
        function emitSeparated(
            orig: Sentence, map: Map<string, SyntaxElement[]>, ref: SyntaxElement | SyntaxElement[])
        {
            const sent = orig.cloneMeta();
            const elem = Array.isArray(ref) ? N.LIST(...ref) : ref;
            console.log('emit: ', NUtils.debugPrintElement(elem));
            sent.elements = [
                NUtils.cloneElement(map.get('subj')![0]),
                N.P(...NUtils.cloneElements(map.get('any')!), elem)
            ];
            newNarrative.push(sent);
        }

        const tempCxt = this.context.clone();
        tempCxt.proximatlySetPronouns.clear();
        const newNarrative: Sentence[] = [];
        const pattern0 = [SP.or('or', SP.signifier('sig'), SP.LIST('list'))];
        const pattern1 = [SP.S('subj'), SP.P().with(
            SP.any('any').asterisk(), 
            SP.O().with(SP.LIST('list'))
        )];
        let lastMatch: Map<string, SyntaxElement[]> | null = null;
        for (let i = 0; i < narrative.length; i++) {
            let s = narrative[i];
            // do back split
            if (lastMatch) {
                let matches = SyntaxPattern.matchPart(pattern0, s.elements);
                for (let match of matches) {
                    let removed: Signifier | Signifier[] | null = null;
                    if (match.capture.has('sig')) {
                        let sig = match.capture.get('sig')![0] as Signifier;
                        removed = removeIfInList(
                            lastMatch.get('list')![0] as SyntaxNode, sig.object);
                    } else {
                        let list = match.capture.get('list')![0] as SyntaxNode;
                        removed = removeIfInList(
                            lastMatch.get('list')![0] as SyntaxNode, 
                            list.children.map(
                                (x) => x instanceof Signifier ? x.object : []).flat());
                    }
                    if (removed) {
                        emitSeparated(narrative[i-1], lastMatch, removed);
                        console.log('... from back');
                        break;
                    }
                }
                lastMatch = null;
            }

            // do front split
            const map = SyntaxPattern.match(pattern1, s.elements);
            if (map) {
                lastMatch = map;
                // process the sentence up to the list
                tempCxt.processSentencePronoun(
                    [...map.get('subj')!, ...map.get('any')!]);
                // search for signifiers referable by pronoun
                const node = map.get('list')![0] as SyntaxNode;
                const separateSigs: (Signifier | Signifier[])[] = [];
                let done = false;
                for (const x of PronounList) {
                    let sig = removeIfInList(node, tempCxt.cache[x]);
                    if (sig) {
                        separateSigs.push(sig);
                        if (done) Debug.warn('separated more than once!')
                        done = true;
                    }
                }
                // emit separated sentences
                for (const ref of separateSigs) emitSeparated(s, map, ref);
                if (done) console.log('...from front');
                // if no LIST is needed, just skip pushing this sentence
                if (node.children.length == 0) continue;
                // process the rest
                tempCxt.processElementPronoun(node);
            } else {
                // process the sentence normally
                tempCxt.processSentencePronoun(s.elements);
            }
            newNarrative.push(s);
        }
        // replace LISTs with only one item
        const pattern2 = [SP.LIST().with(SP.any('only'))];
        for (const s of narrative)
            SyntaxPattern.matchPart(pattern2, s.elements, -1,
                (x) => [x.capture.get('only')![0]]);
        return newNarrative;
    }

    /** 
     * Combine simple sentences with same subject to a complex sentence.
     * A拿起X。A拿起Y。=> A拿起X，又拿起Y。
     * A拿起X。A吃了它。A大叫了一声。=> A拿起X，吃了它，然后大叫了一声。
     * A拿着X。A搬着Y。=> A拿着X，[还]搬着Y。
     */
    #Opt4(narrative: Sentence[]) {
        // ADV_UNDONE: support stative verbs
        function getComplexSentence(
            subj: SyntaxNode, preds: [SyntaxNode, VerbElement][]): SyntaxElement[]
        {
            Debug.assert(preds.length > 1);
            console.log('combining', preds);
            const list: SyntaxNode[][] = [];
            for (let i = 0; i < preds.length;) {
                const v0 = preds[i][1];
                const temp: SyntaxNode[] = [];
                for (; i < preds.length; i++) {
                    if (preds[i][1].equals(v0))
                        temp.push(preds[i][0]);
                    else break;
                }
                if (temp.length > 1)
                    temp[temp.length - 1].children.unshift(
                        v0.isStative ? '还' : '又');
                list.push(temp);
            }
            const list2: SyntaxElement[][] = [];
            for (let i = 0; i < list.length;) {
                if (list[i].length > 1) {
                    list2.push(list[i]);
                    i++;
                    continue;
                }
                const temp: SyntaxElement[] = [];
                for (; i < list.length; i++)
                    if (list[i].length == 1)
                        temp.push(...list[i]);
                    else break;
                list2.push(temp);
            }
            if (list.length > 1)
                list[list.length - 1][0].children.unshift(
                    preds[0][1].isStative ? '' : '然后');
            return [subj, N.COORD_PREDS(...list2.flat())];
        }

        const pattern0 = [
            SP.S('subj'), 
            SP.P('pred').with(SP.any().asterisk(), SP.V('v'), SP.any().asterisk())];
        const newNarrative: Sentence[] = [];
        let i = 0;
        while (i < narrative.length) {
            const s0 = narrative[i];
            const preds: [SyntaxNode, VerbElement][] = [];
            const usedVerb = new Set<string>;
            let subj: SyntaxNode | null = null;
            let lastVerb = '';
            let end = 0;
            let j = i;
            for (; j < narrative.length; j++) {
                let s = narrative[j];
                //if (s.group != s0.group) break;
                let map = SyntaxPattern.match(pattern0, s.elements);
                if (!map) break;
                // console.log('match opt4:', ...s.elements.map((x) => NUtils.debugPrintElement(x)));
                if (!subj) subj = map.get('subj')![0] as SyntaxNode;
                else if (!NUtils.equals(map.get('subj')![0], subj)) break;
                const verb = map.get('v')![0] as VerbElement;
                // break up the complex sentence if a verb has recurred and is not
                // adjacent to its previous occurrence
                if (usedVerb.has(verb.form) && verb.form != lastVerb) break;
                usedVerb.add(verb.form);
                lastVerb = verb.form;
                preds.push([map.get('pred')![0] as SyntaxNode, verb]);
            }
            if (preds.length < 2) {
                j = Math.max(j, i+1);
                newNarrative.push(...narrative.slice(i, j));
            } else {
                Debug.assert(subj !== null);
                let sent = new Sentence(getComplexSentence(subj, preds));
                sent.group = s0.group;
                newNarrative.push(sent);
            }
            i = j;
        }
        return newNarrative;
    }

    updatePronouns(narrative: Sentence[]) {
        this.context.proximatlySetPronouns.clear();
        for (const s of narrative)
            s.elements = this.context.processSentencePronoun(s.elements);
    }

    /**
     * - Removes 'implicit' and 'thought' sentences when necessary.
     * - Collapses same subjects.
     *    - A拿起X。A拿起Y。=> A拿起X和Y。(same verb)
     *    - A拿起X。A吃掉了Y。=> A拿起X，然后吃掉了Y。(different verb)
     *    - A拿起X。A拿不起来Y。=> A拿起X，但是拿不起来Y。(with turning conjunction)
     * - Collapses same objects.
     *    - A摸了摸X。B摸了摸X。=> A和B摸了摸X。
     * - 
     */
    optimize(narrative: Sentence[]): Sentence[] {
        if (narrative.length < 1) return narrative;
        narrative = Narrative.#Opt1.process(narrative);
        narrative = Narrative.#Opt2.process(narrative);
        narrative = this.#Opt3(narrative);
        narrative = this.#Opt4(narrative);
        return narrative;
    }

    // ADV_TODO: support callback here?
    #confabulateList(elem: SyntaxNode) {
        function hasOtherInstancesAround(obj: AObject) {
            if (obj.prototype === null) return false;
            let list: AObject[] = [];
            if (obj.parentPart) list = [...obj.parentPart.parts];
            else if (obj.parentContainer) {
                if (obj.parentContainer instanceof AObject) 
                    list = [...obj.parentContainer.contents];
                else if (obj.parentContainer instanceof AZone)
                    list = [...obj.parentContainer.objects];
                else
                    list = [...obj.parentContainer.inventory];
            } else return false;
            return list.some((x) => x.prototype == obj.prototype);
        }

        const completeSigs = elem.children.filter(
            (x) => x instanceof Signifier) as Signifier[];
        if (completeSigs.length != elem.children.length)
            Debug.warn('NP-list only allows signifiers');
        const objs = completeSigs.filter((x) => x.object instanceof AObject);
        const agents = completeSigs.filter((x) => x.object instanceof Agent);

        const singlePart: [AObject | null, Signifier[]][] = [];
        const pluralPart: [AObject | null, Signifier[]][] = [];
        [...U.groupPrototypeSignifier(objs).entries()].forEach((x) =>
            (x[1].length == 1 ? singlePart : pluralPart).push(x));
        const results: string[] = [];
        for (const [proto, sigs] of [...singlePart, ...pluralPart]) {
            if (proto === null) {
                results.push(...sigs.map((sig) => {
                    Debug.assert(sig.object instanceof AObject);
                    if (sig.forceIndefiniteArticle) return '一' + (sig.object.class?.determiner ?? '个') + sig.object.name.value;
                    else return sig.object.name.value;
                }));
            } else {
                let count = '';
                let first = sigs[0].object;
                Debug.assert(first instanceof AObject);
                if (sigs.length > 1 || sigs[0].forceIndefiniteArticle 
                    || hasOtherInstancesAround(first))
                        count = NAlgorithms.expressNumber(objs.length) 
                            + (proto.class?.determiner ?? '个');
                results.push(count + proto.name.value)
            }
        }

        for (const ag of agents) {
            Debug.assert(ag.object instanceof Agent);
            results.push(ag.object.name);
        }

        let result = '';
        for (let i = 0; i < results.length; i++) {
            result += results[i];
            if (i < results.length - 2) result += '、';
            else if (i == results.length - 2) result += '和';
        }
        return result;
    }

    #confabulateElement(elem: SyntaxElement, callback?: ConfabulationCallback): string {
        let value = '<?>';
        if (typeof elem == 'string') {
            value = elem;
        } else if (elem instanceof VerbElement) {
            if (elem.isPerfective)
                value = `${elem.form}了` + (elem.form.length == 1 ? `${elem.form}` : '');
            else value = elem.form;
        } else if (elem instanceof Signifier) {
            if (elem.form != '') return elem.form;
            // ADV_TODO
            if (elem.object instanceof AObject) {
                value = '';
                if (elem.forceIndefiniteArticle)
                    value = '一' + (elem.object.class?.determiner ?? '个');
                value += elem.object.name.value;
            } else {
                value = elem.object.name;
            }
        } else {
            Debug.assert(NUtils.isNode(elem));
            if (elem.type == 'nounPhraseList') {
                value = this.#confabulateList(elem);
            } else {
                const list = elem.children.map(
                    (x) => this.#confabulateElement(x, callback));
                if (elem.type == 'coordinatingPredicatives')
                    value = list.join('，');
                else
                    value = list.join('');
            }
        }
        if (callback)
            value = callback(elem, value) ?? value;
        return value;
    }

    /** 'Confabulate', 'enunciate', 'asseverate', or rather 'ejaculate'.
     */
    confabulate(sent: Sentence, callback?: ConfabulationCallback): string {
        if (sent.elements.length == 0) return '';
        let end = '';
        if (sent.elements.length > 1 || (!NUtils.isNode(sent.elements[0]) || sent.elements[0].type != 'text') && !sent.flags.has(SentenceFlags.Unfinished)) end = '。';
        return sent.elements.map(
            (x) => this.#confabulateElement(x, callback)).join('') + end;
    }
}